import torch
import torch.optim as optim

from util.seed import set_seed
from raw.get_model import get_model

#训练并保存网络
def main(name):
    # 设置随机数种子
    set_seed()

    #创建网络
    model = get_model(name)
    #获取optimizer
    opt = model.get_optimizer()

    batch_size = model.training_params['batch_size']
    #加载训练数据
    train_loader, valid_loader = model.get_train_valid_loader(batch_size = batch_size,
                                                                valid_size = 0.1,
                                                                transform_train = True,
                                                                shuffle = True)
    #加载测试数据
    test_loader = model.get_test_loader(batch_size = batch_size, shuffle = False)

    # 进行训练
    best_ratio = 0 # 最优准确率
    best_state_dict = None # 最优网络参数
    for epoch in range(model.training_params['num_epochs']):
        model.train_once(train_loader, opt, epoch) # 训练一轮
        ratio = model.test(valid_loader) # 测试
        model.adjust_optimizer(opt, epoch) # 调整optimizer参数
        if ratio > best_ratio:
            best_ratio = ratio
            best_state_dict = model.state_dict_clone()
            print('Validate: Accuracy increased')
        else:
            print('Validate: Accuracy decreased')
        print()

    # 进行测试
    model.load_state_dict(best_state_dict)
    model.test(test_loader)

    # 保存最优网络参数
    path = 'data/raw_models/' + name + '.pt'
    torch.save(best_state_dict, path)
    print('Model saved to ' + path)
